package fun.ticsmyc.graph;

import org.junit.Test;

import java.util.Arrays;
import java.util.HashSet;
import java.util.PriorityQueue;
import java.util.Set;

/**
 * 最小生成树： Prim算法
 * leetcode 1584 连接所有点的最小费用
 * @author Ticsmyc
 * @date 2021-01-19 14:04
 */
public class Prim {

    /**
     * 生成邻接表
     * @param points
     * @return
     */
    public int[][] createGraph(int[][] points){
        int n = points.length;
        int[][] graph = new int[n][n];
        for(int i=0;i<n;++i){
            for(int j=i+1;j<n;++j){
                graph[i][j] = dis(points[i],points[j]);
                graph[j][i] = graph[i][j];
            }
        }
        return graph;
    }

    /**
     * 邻接矩阵的prim
     * @param graph
     * @return
     */
    public int prim(int[][] graph){
        int n  = graph.length;
        int ans =0;
        Set<Integer> visited = new HashSet<>();
        int[] dis = new int[n]; //统计 U-V集合间的距离

        //以0点作为起点 ,加入V集合
        for(int i=0;i<n;++i){
            dis[i] = graph[0][i];
        }
        visited.add(0);

        //剩下的n-1个节点
        for(int i=0;i<n-1;++i){
            int distance =Integer.MAX_VALUE;
            int end =0;
            for(int j=0;j<n;++j){
                if(!visited.contains(j) && dis[j] < distance){
                    distance = dis[j];
                    end = j;
                }
            }
            ans += distance;
            visited.add(end);
            //end 即为本次加入生成树的节点
            //更新加入end后，其他节点到生成树的距离
            for(int j=0;j<n;++j){
                if(graph[end][j] < dis[j]){
                    dis[j] = graph[end][j];
                }
            }
        }
        return ans;
    }

    /**
     * 使用优先队列优化的 邻接矩阵的 prim
     * @param graph
     * @return
     */
    public int prim2(int[][] graph){
        int n  =graph.length;
        int ans =0;
        Set<Integer> visited = new HashSet<>();
        int[] dis = new int[n]; //统计 U-V集合间的距离
        Arrays.fill(dis,Integer.MAX_VALUE);
        PriorityQueue<int[]> queue = new PriorityQueue<>((a,b)->a[0]-b[0]);

        queue.offer(new int[]{0,0});
        while(!queue.isEmpty()){
            int[] tmp = queue.poll();
            int d = tmp[0]; //距离
            int index = tmp[1]; //点编号
            if(visited.contains(index)){
                continue;
            }
            visited.add(index);
            ans += d;

            for(int j=0;j<n;++j){
                int distance = graph[index][j];
                if(!visited.contains(j) && dis[j] > distance){
                    dis[j] = distance;
                    queue.offer(new int[]{distance,j});
                }
            }
        }
        return ans;
    }

    public int dis (int[] a, int[] b){
        return Math.abs(a[0]-b[0]) + Math.abs(a[1]-b[1]);
    }

    @Test
    public void test(){
        int[][] points = new int[][]{{0, 0}, {1, 1}, {1, 0}, {-1, 1}};
        int n = points.length;
        int[][] grade = createGraph(points);
        int res = prim(grade);
        System.out.println(res);

    }
}
